'''
@env        : python3.7
@time       : 2020.11.24 ~
@author     : ZhangZhuoli
@func       : RandomForest
@reference  : 
'''
import sys
import pandas as pd 
import numpy as np
import random
import math
import re
if sys.version > '3':
    from functools import reduce

class Tree(object):
    """
    定义决策树
    """
    def __init__(self):
        self.P = 1
        self.N = 0
        self.target_attribute = None
        self.DTree = None
        
    def ispn(self,index):
        '''判断正例反例并返回主要判断'''
        # 初始化
        p = 0
        n = 0
        majority = "Unknow"
        for i in index:
            if examples[i][self.target_attribute] == self.P:
                p += 1
            else:
                n += 1
        if p > n : majority = self.P
        elif n > p : majority = self.N
        return p, n, majority

    def entropy(self,index):
        '''计算熵值'''
        p, n, _ = self.ispn(index)
        p = float(p)
        n = float(n)
        if p*n != 0 :
            return -p/len(index) * np.log2(p/len(index)) - n/len(index) * np.log2(n/len(index))
        else:
            return 0

    def subsetIndex(self,index,key,subset):
        '''索引subset在key中的下标'''
        return list(filter(lambda x: examples[x][key] == subset, index))

    def gain(self,index,attr):
        '''计算信息熵'''
        e = self.entropy(index)
        for value in attri_values_table[attr]:
            index_value = self.subsetIndex(index,attr,value)
            e -= float(len(index_value))/len(index) * self.entropy(index_value)
        return e

    def isRoot(self,index,attributes):
        '''通过最大信息熵选择最优特征作为根节点'''
        max_gain = -1
        best_attr = None
        for attr in attributes:
            a_gain = self.gain(index, attr)
            if a_gain > max_gain:
                best_attr = attr
                max_gain = a_gain
        return best_attr


    def id3(self,index,attributes):
        '''ID3算法参数说明:
            index: 带分叉的例集索引；
            attributes: 除目标属性外供决策树训练的属性列表
        '''
        p, n, majority = self.ispn(index)
        if p == 0 : return {'label':self.N, 'p/n':'[' + str(p) + '+,' + str(n) + '-]'}
        if n == 0 : return {'label':self.P, 'p/n':'[' + str(p) + '+,' + str(n) + '-]'}
        if len(attributes) == 0: return {'label':majority, 'p/n':'['+str(p)+'+,'+str(n)+'-]'}

        Tree_record = {'decision':self.isRoot(index,attributes),'p/n':'['+str(p)+'+,'+str(n)+'-]'}
        for value in attri_values_table[Tree_record['decision']]:
            index_v = self.subsetIndex(index,Tree_record['decision'],value)
            if len(index_v) == 0:
                Tree_record[str(value)] = {'label':majority, 'p/n':'[0+,0-]'}
            else:
                attributes_v = list(attributes)
                attributes_v.remove(Tree_record['decision'])
                Tree_record[str(value)] = self.id3(index_v, attributes_v)
        return Tree_record
        # np.save('Tree_record.npy',Tree_record)
    
    def treeview(self, tree, space = ''):
        """
        绘制树状图
        """
        if not "label" in tree:
            print(space + '|--'+'decision : ' + str(tree['decision'])+ tree['p/n'])
            for value in attri_values_table[tree['decision']]:
                print(space + '|--' + str(value))
                self.treeview(tree[str(value)],space = ':       ')
        else:
            print(space + 'label : '+ str(tree['label']) + tree['p/n'])


if __name__ == "__main__":

    data_train = pd.read_csv('./data/train.csv')
    with open('./data/train.csv','r') as f:
        # 读取完整数据
        examples = f.readlines()
        examples = list(map(lambda s: re.split(',',s.rstrip('\n')),examples[:]))
        
        # 读取固定行数数据
        # examples = []
        # for i in range(500):
        #     examples.append(re.split(',',f.readline().rstrip('\n')))
    # print(examples)

    names = examples.pop(0)
    attributes = names[1:-1]
    target_attribute = names[-1]

    """ for i in range(len(examples)):
        for j in range(len(examples[0])):
            examples[i][j] = int(examples[i][j]) """
    # 列表推导式转数据类型为int
    examples = [ dict(zip(names[1:],[int(examples[i][j]) for j in range(1,len(examples[0]))])) for i in range(len(examples)) ]
    # print('examples:',examples)

    '''构造特征值表'''
    # 导入提前处理并保存的文件
    attri_values_table = np.load('./Attri_Values.npy',allow_pickle=True).item()
    # print(attri_values_table)

    """ attri_values_table = []
    for j in range(len(examples[0])):
        temp = []
        for i in range(len(examples)):
            if examples[i][j] not in temp:
                temp.append(examples[i][j])
        attri_values_table.append(temp)
    attri_values_table = dict(zip(names[1:],attri_values_table))
    print(attri_values_table) """

    DecisionTree = Tree()
    DecisionTree.target_attribute = target_attribute
    Dtree = DecisionTree.id3(range(len(examples)),attributes)
    DecisionTree.DTree = Dtree
    np.save('Tree_record.npy',Dtree)
    DecisionTree.treeview(DecisionTree.DTree)